import sys
sys.path.append("D:\wordSpace\langchainQA")

from BCEmbedding import RerankerModel
from typing import List, Optional, Sequence, Union,Any
from langchain.retrievers.document_compressors.base import BaseDocumentCompressor

from langchain_core.documents import  Document
from langchain.callbacks.manager import Callbacks
from app.embedding.base import baseReRank
from app.config.model_config import  RERANKER_MODEL,EMBEDDING_DEVICE_MAP
from app.utils.singleton import SignletonFunc,Singleton
from app.config.log_config import loggerqa

import time
@SignletonFunc
class MyDocumentCompressor(BaseDocumentCompressor):
    top_n = 3
    
    # def __init__(self,path=RERANKER_MODEL,device=EMBEDDING_DEVICE_MAP['ranker_model']):
    #     self.model = RerankerModel(model_name_or_path=path,device=device)
    #     super().__init__()
    
    # class Config:
    #         """Configuration for this pydantic object."""

    #         model:Optional[RerankerModel]=None
    def  compress_documents(self,
        documents: Sequence[Document],
        query: str,callbacks: Optional[Callbacks] = None):
        s=time.time()
        passages=[doc.page_content for doc in documents]
        model = baseReRank
        rerank_response=model.rerank(query,passages)
        rerank_passages=rerank_response["rerank_passages"]
        rerank_scores=rerank_response["rerank_scores"]
        rerank_ids=rerank_response["rerank_ids"][:self.top_n]
        loggerqa.info(f"重排耗时:{time.time()-s}")
        return [Document(page_content=rerank_passages[i],metadata={"relevance_score":rerank_scores[i],"id":i})  for i in rerank_ids]
    async def  acompress_documents(self,
        documents: Sequence[Document],
        query: str,callbacks: Optional[Callbacks] = None):
        s=time.time()
        passages=[doc.page_content for doc in documents]
        model = baseReRank
        rerank_response=await model.arerank(query,passages)
        rerank_passages=rerank_response["rerank_passages"]
        rerank_scores=rerank_response["rerank_scores"]
        rerank_ids=rerank_response["rerank_ids"][:self.top_n]
        loggerqa.info(f"异步重排耗时:{time.time()-s}")
        return [Document(page_content=rerank_passages[i],metadata={"relevance_score":rerank_scores[i],"id":i})  for i in rerank_ids]
if __name__ == '__main__':
    import asyncio
    import time
    app=MyDocumentCompressor()
    async def test():
        
        s=time.time()
        res=app.compress_documents(documents=[Document(page_content="hello world"),Document(page_content="你好")],query="hello world")
        res=await app.acompress_documents(documents=[Document(page_content="hello world"),Document(page_content="你好")],query="hello world")
    async def main():
        s=time.time()
        task=[test() for i in range(100)]
        await asyncio.gather(*task)
        print(time.time()-s,'====')
    asyncio.run(main())
    